package com.xinxin.leetcode.problem347;

import java.util.*;

/**
 * @author 史鑫鑫
 * @date Created in 2019/5/28 4:26
 */
class Solution {
    private class Freq implements Comparable<Freq> {
        int e, freq;

        Freq(int e, int freq) {
            this.e = e;
            this.freq = freq;
        }

        @Override
        public int compareTo(Freq another) {
            return Integer.compare(this.freq, another.freq);
        }
    }

    public List<Integer> topKFrequent(int[] nums, int k) {
        TreeMap<Integer, Integer> map = new TreeMap<>();
        for (int num : nums) {
            if (map.containsKey(num)) {
                map.put(num, map.get(num) + 1);
            } else {
                map.put(num, 1);
            }
        }
        System.out.println(map);
        PriorityQueue<Freq> pq = new PriorityQueue<Freq>();
        for (int key : map.keySet()) {
            if (pq.size() < k) {
                pq.add(new Freq(key, map.get(key)));
            } else if (map.get(key) > pq.peek().freq) {
                pq.poll();
                pq.add(new Freq(key, map.get(key)));
            }
        }
        List<Integer> res = new LinkedList<>();
        while (!pq.isEmpty()) {
            res.add(pq.poll().e);
        }
        return res;
    }

    public static void main(String[] args) {
        int[] nums = {4, 1, -1, 2, -1, 2, 3};
        int k = 2;
        List list = new Solution().topKFrequent(nums, k);
        System.out.println(list);
    }
}
